"""
Case Type   : 多块预读
Case Name   : 开启多块预读，存储过程/函数中扫描表
Create At   : 2024/12/03
Owner       : @li-xin12345
Description :
    1、建表
    2、创建存储过程
    3、创建自定义函数
    4、调用存储过程
    5、调用自定义函数并获取结果
    6、查看日志中打印的预读块数量
    7、清理环境
Expect      :
    1、成功
    2、成功
    3、成功
    4、成功
    5、成功
    6、预读了设置的块数量
    7、成功
History     :
"""
import os
import unittest

from yat.test import Node
from yat.test import macro

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger


class SessionTest(unittest.TestCase):
    def setUp(self):
        self.log = Logger()
        self.log.info(f'-----{os.path.basename(__file__)} start-----')
        self.constant = Constant()
        self.common = Common()
        self.pri_sh = CommonSH('PrimaryDbUser')
        self.pri_user_node = Node('PrimaryDbUser')
        self.tb_name = 't_multi_bulk_read_0011'
        self.func_name = 'f_multi_bulk_read_0011'
        self.proc_name = 'p_multi_bulk_read_0011'
        self.pg_log = \
            os.path.join(macro.PG_LOG_PATH, macro.DN_NODE_NAME.split('/')[0])

    def test_multi_bulk_read(self):
        text = '----step1: 建表；expect: 成功----'
        self.log.info(text)
        sql = f'''drop table if exists {self.tb_name}_1, 
        {self.tb_name}_2 cascade;
        create table {self.tb_name}_1 (
            c_int int,
            c_char char(10),
            c_bool bool default false
        );
        insert into {self.tb_name}_1 values(
        generate_series(1,100000), 'char'||generate_series(1,100000));
        create table {self.tb_name}_2 as select * from {self.tb_name}_1;
        '''
        res = self.pri_sh.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.DROP_TABLE_SUCCESS, res, f'执行失败：{text}')
        self.assertIn(self.constant.CREATE_TABLE_SUCCESS,
                      res, f'执行失败：{text}')
        self.assertEqual(res.count('INSERT 0 100000'), 2, f'执行失败：{text}')

        text = '----step2: 创建存储过程；expect: 成功----'
        self.log.info(text)
        sql = f'''drop procedure if exists {self.proc_name};
        create procedure {self.proc_name}(tb_name varchar) is
        begin
            execute immediate 'select * from ' || tb_name || ';';
            return;
        end;
        '''
        res = self.pri_sh.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.DROP_PROCEDURE_SUCCESS_MSG, res,
                      f'执行失败：{text}')
        self.assertIn(self.constant.CREATE_PROCEDURE_SUCCESS_MSG, res,
                      f'执行失败：{text}')

        text = '----step3: 创建自定义函数；expect: 成功----'
        self.log.info(text)
        sql = f'''drop function if exists {self.func_name};
        create function {self.func_name}() 
        returns table (c1 int, c2 char(10), c3 bool) as \$\$
        begin
            return query select * from {self.tb_name}_2 order by 1 desc 
            limit 10;
        end;
        \$\$ language plpgsql;
        '''
        res = self.pri_sh.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, res,
                      f'执行失败：{text}')
        self.assertIn(self.constant.CREATE_FUNCTION_SUCCESS_MSG, res,
                      f'执行失败：{text}')

        msg = self.pri_sh.restart_db_cluster()
        self.log.info(msg)
        status = self.pri_sh.get_db_cluster_status()
        self.assertTrue("Degraded" in status or "Normal" in status)

        text = '----step4: 调用存储过程；expect: 预读了设置的块数量----'
        self.log.info(text)
        sql = f"set heap_bulk_read_size = 16; " \
              f"call {self.proc_name}('{self.tb_name}_1');"
        res = self.pri_sh.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.SET_SUCCESS_MSG, res, f'执行失败：{text}')
        self.assertIn('1 row', res, f'执行失败：{text}')

        text = '----step5: 调用自定义函数并获取结果；expect: 预读了设置的块数量----'
        self.log.info(text)
        sql = f"set heap_bulk_read_size = 16; " \
              f"select * from {self.func_name}(); "
        res = self.pri_sh.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.SET_SUCCESS_MSG, res, f'执行失败：{text}')
        self.assertIn('10 rows', res, f'执行失败：{text}')

        text = '----step6: 查看日志中打印的预读块数量；expect: 预读了设置的块数量----'
        self.log.info(text)
        cmd = f"cd {self.pg_log};" \
              f"ls -t | head -1 | awk '{{{{print $1}}}}' | " \
              f"xargs cat | tail -30 | " \
              f"grep 'End of pre-Read, the max blocks batch is'"
        self.log.info(cmd)
        res = self.common.get_sh_result(self.pri_user_node, cmd)
        self.log.info(res)
        self.assertIsNotNone(res, f'执行失败：{text}')

    def tearDown(self):
        text = '----step7: 清理环境;expect:成功----'
        self.log.info(text)
        sql = f"drop table if exists {self.tb_name}_1, " \
              f"{self.tb_name}_2 cascade;" \
              f"drop function if exists {self.func_name};" \
              f"drop procedure if exists {self.proc_name};"
        res = self.pri_sh.execut_db_sql(sql)
        self.log.info(res)
        self.assertIn(self.constant.DROP_TABLE_SUCCESS, res, f'执行失败：{text}')
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, res,
                      f'执行失败：{text}')
        self.assertIn(self.constant.DROP_PROCEDURE_SUCCESS_MSG, res,
                      f'执行失败：{text}')
        self.log.info(f'-----{os.path.basename(__file__)} end-----')
